from cion.gates import *
#from AD9910 import AD9910


virtual_chapter_dict = {
    ('ETTIT',5):'00000001 00000001 00000001', 
    'EIT':'00000001 00000001 00000001', 
    'Doppler':'00000000 00000000 00001111', 
    #'Detection':'10000001 00000000 00001111',
    'Detection': '10000001 00000000 00001111,[10000000 00000000 00000000, 2],[10000000 00000000 00010000,2]',
    'optical_pumping':'00000100 00000100 00000100',
    'RSB':'00001000 00001000 00001000', 
    'Zero': '10000000 00000000 00000000', 
    'Full': '11111111 11111111 11111111',
    'Rx': '10010010 00010000 00000000',
    'parallel_xx':'00110000 00000000 00000000',
    'parallel_yy':'00000000 00110000 00000000', 
    'parallel_zz':'00000000 00000000 00110000', 
    'global_xx':'11000000 00000000 00000000',
    'global_xx':'00000000 11000000 00000000',
    'global_xx':'00000000 00000000 11000000'
              }

#AWG pulse种类和FPGA完全一样
#输入参数t是距离trigger的绝对时间
#FPGA模块和AWG模块时序同步
#seq类需要灵活地修改某些AWG pulses的参数，建议用String label的方式

#下面是一个分段函数的例子

'''
0.45*Sin(2.03*t+0.48)+0.45*Sin(-2.03*t-0.48), t<9.33
0.45*Sin(2.03*t+0.308)+0.45*Sin(-2.03*t-0.308), 9.33<t<18.66
0.45*Sin(2.03*t+0.086)+0.45*Sin(-2.03*t-0.086), 18.66<t<27.99
0.45*Sin(2.03*t-0.086)+0.45*Sin(-2.03*t+0.086), 27.99<t<37.32
0.45*Sin(2.03*t+0.308)+0.45*Sin(-2.03*t-0.308), 37.32<t<46.65
0.45*Sin(2.03*t+0.48)+0.45*Sin(-2.03*t-0.48), 46.65<t<55.98
'''

class Doppler(BaseGate):
    hardware_amp, hardware_freq, hardware_phase = None, None, None
    def __init__(self, duration, latency = 0, awg_flag = None, amp = None, freq = None, phase = None, label = "Doppler"):
        super().__init__(duration, latency, awg_flag, 'Doppler', amp, freq, phase, label)
        self.hardware_update_function = self.update_function_define
    def direct_update_hardware(self):
        # the first two parameters representing "device number" and "channel", which could be modified
        assert Doppler.hardware_amp != None and Doppler.hardware_freq != None and Doppler.hardware_phase != None
        hardware.AD9910.WriteProf(1, 3, Doppler.hardware_amp, Doppler.hardware_freq, Doppler.hardware_phase)
        #return AD9910.WriteProf(1, 3, hardware_amp, hardware_freq, hardware_phase)
    def update_function_define(self, amp, freq, phase):
        return AD9910.WriteProf(1, 3, amp, freq, phase)
    def update_hardware_parameters(self, amp, freq, phase):
        Doppler.hardware_amp = amp
        Doppler.hardware_freq = freq
        Doppler.hardware_phase = phase

class EIT(BaseGate):
    hardware_amp, hardware_freq, hardware_phase = None, None, None
    def __init__(self, duration, latency = 0, awg_flag = None, amp = None, freq = None, phase = None, label = "EIT"):
        super().__init__(duration, latency, awg_flag, 'EIT', amp, freq, phase, label)
    def direct_update_hardware(self):
        assert EIT.hardware_amp != None and EIT.hardware_freq != None and EIT.hardware_phase != None
        pass#return AD9910.WriteProf(1, 3, hardware_amp, hardware_freq, hardware_phase)
    '''
    def __init__(self, duration, amp = 20, freq = 100, phase = 0, awg = None, fpga = None):
        self.duration = duration
        self.pulse_type = 'EIT'
    '''

class Detection(BaseGate):
    hardware_amp, hardware_freq, hardware_phase = None, None, None
    def __init__(self, duration, latency = 0, awg_flag = None, amp = None, freq = None, phase = None, label = "Detection"):
        super().__init__(duration, latency, awg_flag, 'Detection', amp, freq, phase, label)

    def direct_update_hardware(self):
        assert Detection.hardware_amp != None and Detection.hardware_freq != None and Detection.hardware_phase != None
        pass#return AD9910.WriteProf(1, 3, hardware_amp, hardware_freq, hardware_phase)

class Optical_pumping(BaseGate):
    hardware_amp, hardware_freq, hardware_phase = None, None, None
    def __init__(self, duration, latency = 0, awg_flag = None, amp = None, freq = None, phase = None, label = "Optical_pumping"):
        super().__init__(duration, latency, awg_flag, 'optical_pumping', amp, freq, phase, label)

    def direct_update_hardware(self):
        assert Optical_pumping.hardware_amp != None and Optical_pumping.hardware_freq != None and Optical_pumping.hardware_phase != None
        pass#return AD9910.WriteProf(1, 3, hardware_amp, hardware_freq, hardware_phase)

class RSB(BaseGate):
    hardware_amp, hardware_freq, hardware_phase = None, None, None
    def __init__(self, duration, latency = 0, awg_flag = None, amp = None, freq = None, phase = None, label = "RSB"):
        super().__init__(duration, latency, awg_flag, 'RSB', amp, freq, phase, label)

    def direct_update_hardware(self):
        assert RSB.hardware_amp != None and RSB.hardware_freq != None and RSB.hardware_phase != None
        pass#return AD9910.WriteProf(1, 3, hardware_amp, hardware_freq, hardware_phase)

class Zero(BaseGate):
    hardware_amp, hardware_freq, hardware_phase = None, None, None
    def __init__(self, duration, latency = 0, awg_flag = None, amp = None, freq = None, phase = None, label = "Zero"):
        super().__init__(duration, latency, awg_flag, 'Zero', amp, freq, phase, label)
        
class Full(BaseGate):
    hardware_amp, hardware_freq, hardware_phase = None, None, None
    def __init__(self, duration, latency = 0, awg_flag = None, amp = None, freq = None, phase = None, label = "Full"):
        super().__init__(duration, latency, awg_flag, 'Full', amp, freq, phase, label)


class AWGBaseGate(BaseGate):
    def __init__(self, duration, latency = 0, awg_flag = None, amp = None, freq = None, phase = None,\
                 label = None, segment_number = -1, time_intervals = None):
        super().__init__(duration, latency, awg_flag, 'Rx', amp, freq, phase, label)
        self.segment_flag = False
        self.time_intervals = None
        if segment_number > 0:
            self.segment_flag = True
            self.segment_number = segment_number
            self.time_intervals = time_intervals
            assert isinstance(self.time_intervals, list)
            #print(self.time_intervals)
            for pair in self.time_intervals:
                #print(pair, len(pair))
                assert len(pair) == 2
            assert len(self.time_intervals) == self.segment_number

class Rx(BaseGate):
    # specify the following arguments
    # amp, omega, duration
    def __init__(self, duration, latency = 0, awg_flag = None, amp = None, freq = None, phase = None,\
                 label = None, segment_number = -1, time_intervals = None):
        super().__init__(duration, latency, awg_flag, 'Rx', amp, freq, phase, label)
        #确定该函数的分段数
        #fix_num = 4
        #self.segment_num = fix_num
        self.segment_flag = False
        self.time_intervals = None
        if segment_number > 0:
            self.segment_flag = True
            self.segment_number = segment_number
            self.time_intervals = time_intervals
            assert isinstance(self.time_intervals, list)
            for pair in self.time_intervals:
                assert len(pair) == 2
            assert len(self.time_intervals) == self.segment_number
            
    def update_hardware(self):
        None
        
    
    def reset_hardware(self):
        None

    '''
    0.45*Sin(2.03*t+0.48)+0.45*Sin(-2.03*(t+t0)-0.48), t<9.33
    0.45*Sin(2.03*t+0.308)+0.45*Sin(-2.03*t-0.308), 9.33<t<18.66
    0.45*Sin(2.03*t+0.086)+0.45*Sin(-2.03*t-0.086), 18.66<t<27.99
    0.45*Sin(2.03*t-0.086)+0.45*Sin(-2.03*t+0.086), 27.99<t<37.32
    0.45*Sin(2.03*t+0.308)+0.45*Sin(-2.03*t-0.308), 37.32<t<46.65
    0.45*Sin(2.03*t+0.48)+0.45*Sin(-2.03*t-0.48), 46.65<t<55.98
    0, 55.98<t
    '''

    #
    #phase_list = [0.48, 0.308, 0.086]

    #func_list = [func0,func1,func2,func3,func4,func5]

    # make sure that self.duration < 55.98


# class for multi-qubit awg gate, such as 'global entanglement'

# the following contents are for self_defined pulses/gates
# we give a example here:
def example_pulse(ion, amp, freq, phase, duration, pulse_type = None):
    return BaseGate(ion, amp, freq, phase, duration, pulse_type)
# When we call this function 'example(2, 1, 2, np.pi/4, 100)', a BaseGate object will be generated
# Performed on 2-th ion, pulse amplitude = 1, angular frequency is 2(MHZ), phase is pi/4, duration is 100 microseconds.

def carrier_single(phase, pulse_type):
    # specify the following arguments
    # amp, omega, duration
    amp = None
    omega = None
    duration = None
    return BaseGate(amp, omega, phase, duration, pulse_type)

def global_gate(duration, pulse_type = 'xx'):
    amp = None
    freq = None
    phase = None
    pulse_type = 'global_' + pulse_type
    return BaseGate(amp, freq, phase, duration, pulse_type)
